import torch
import cv2 as cv
from model.zhnnet import ZhnNet
from model.decode import decode_box, image_generate_conf, image_generate_loc

print('Detect object.')
classify = False
model = ZhnNet(classify=classify)
model.load_state_dict(torch.load('zhnnet.pth'))
model.eval()
image_origin = cv.imread('E:/dataset/instrument/test3.png')
assert image_origin is not None, 'Image does not exit.'
image = image_origin.transpose(2, 0, 1) / 256
print('Network loading complete.')
with torch.no_grad():
    image = torch.tensor(image, dtype=torch.float32).unsqueeze(0)
    predict = model(image)
    predict = decode_box(predict)
image_origin = image_generate_conf(image_origin, predict) if classify else image_generate_loc(image_origin, predict)
cv.imshow('test', image_origin)
cv.waitKey(0)
